"""
抖音视频下载器 - 验证码处理版
加入手动验证处理步骤
"""

import os
import time
import requests
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from webdriver_manager.chrome import ChromeDriverManager

def download_douyin_video(url, headless=False):
    """
    使用 Selenium 下载抖音视频 - 处理验证码版本
    
    参数:
        url (str): 抖音视频的分享链接
        headless (bool): 是否使用无头模式（不显示浏览器窗口）
    """
    print("=" * 50)
    print(f"开始处理抖音视频: {url}")
    
    # 设置 Chrome 选项
    chrome_options = Options()
    if headless:
        chrome_options.add_argument("--headless")  # 无头模式
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    chrome_options.add_argument("--disable-infobars")
    chrome_options.add_argument("--disable-extensions")
    chrome_options.add_argument("--disable-notifications")
    chrome_options.add_argument("--disable-software-rasterizer")
    chrome_options.add_argument("--disable-webgl")
    chrome_options.add_experimental_option('excludeSwitches', ['enable-logging'])
    
    # 设置用户代理
    user_agent = "Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1"
    chrome_options.add_argument(f"user-agent={user_agent}")
    
    # 设置窗口大小（模拟移动设备）
    chrome_options.add_argument("--window-size=375,812")
    
    driver = None
    try:
        print("初始化浏览器驱动...")
        # 使用 webdriver-manager 自动管理驱动
        service = Service(ChromeDriverManager().install())
        driver = webdriver.Chrome(service=service, options=chrome_options)
        
        print("正在访问页面...")
        driver.get(url)
        
        # 等待页面加载
        print("等待页面加载...")
        time.sleep(3)
        
        # 检查是否需要登录
        if "登录" in driver.page_source:
            print("检测到登录页面，需要手动登录...")
            if headless:
                print("无头模式下无法登录，请使用非无头模式(headless=False)")
                return None
            else:
                print("请在浏览器中完成登录，完成后按回车继续...")
                input()
        
        # 检查是否需要验证
        print("检查是否需要验证...")
        captcha_detected = False
        
        # 检查验证码元素是否存在
        try:
            WebDriverWait(driver, 5).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, ".captcha_verify_container, .captcha-container, .captcha-dialog"))
            )
            captcha_detected = True
        except:
            pass
        
        # 检查验证码关键词
        if not captcha_detected:
            captcha_keywords = ["验证", "安全验证", "拖动滑块", "点选"]
            for keyword in captcha_keywords:
                if keyword in driver.page_source:
                    captcha_detected = True
                    break
        
        # 处理验证码
        if captcha_detected:
            if headless:
                print("检测到验证码，但无头模式下无法处理，请使用非无头模式(headless=False)")
                return None
            else:
                print("=" * 50)
                print("检测到验证码，请手动完成验证！")
                print("1. 在浏览器窗口中完成验证")
                print("2. 验证完成后，请在此窗口按回车继续")
                print("=" * 50)
                input("完成验证后按回车继续...")
                print("继续执行脚本...")
        
        # 获取页面标题
        title = driver.title.replace(" - 抖音", "").strip()
        print(f"视频标题: {title}")
        
        # 尝试点击播放按钮（如果存在）
        try:
            play_button = WebDriverWait(driver, 10).until(
                EC.element_to_be_clickable((By.CSS_SELECTOR, 'div[data-e2e="video-player"] button'))
            )
            play_button.click()
            print("点击了播放按钮")
            time.sleep(2)
        except:
            print("未找到播放按钮或无需点击")
        
        # 查找视频元素 - 使用更可靠的方法
        print("查找视频元素...")
        
        # 方法1: 使用JavaScript获取视频元素
        video_elements = driver.execute_script("""
            // 尝试获取所有video元素
            let videos = document.querySelectorAll('video');
            if (videos.length > 0) return Array.from(videos).map(v => v.src);
            
            // 尝试在Shadow DOM中查找
            let shadowRoots = [];
            document.querySelectorAll('*').forEach(el => {
                if (el.shadowRoot) shadowRoots.push(el.shadowRoot);
            });
            
            for (let root of shadowRoots) {
                videos = root.querySelectorAll('video');
                if (videos.length > 0) return Array.from(videos).map(v => v.src);
            }
            
            return [];
        """)
        
        # 方法2: 如果JS方法失败，使用传统方法
        if not video_elements:
            print("使用传统方法查找视频元素...")
            video_elements = []
            elements = driver.find_elements(By.TAG_NAME, 'video')
            for el in elements:
                src = el.get_attribute('src')
                if src:
                    video_elements.append(src)
        
        if not video_elements:
            print("无法找到视频元素，保存页面截图...")
            driver.save_screenshot("debug.png")
            print("请查看 debug.png 文件进行故障排除")
            return None
        
        print(f"找到 {len(video_elements)} 个可能的视频源")
        
        # 获取视频URL
        video_url = None
        for src in video_elements:
            if src and "http" in src:
                print(f"候选视频URL: {src}")
                # 检查URL是否有效
                if any(key in src for key in [".mp4", ".m3u8", "video/"]):
                    video_url = src
                    print("使用此视频URL")
                    break
        
        if not video_url and video_elements:
            # 如果没有明确的视频URL，选择第一个候选
            video_url = video_elements[0]
            print(f"使用第一个候选视频URL: {video_url}")
        
        if not video_url:
            print("无法获取有效的视频URL")
            return None
        
        print(f"最终视频URL: {video_url}")
        
        # 下载视频
        headers = {
            "User-Agent": user_agent,
            "Referer": driver.current_url,
            "Accept": "*/*",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
            "Range": "bytes=0-"
        }
        
        print("开始下载视频...")
        start_time = time.time()
        
        # 使用requests下载视频
        response = requests.get(video_url, headers=headers, stream=True)
        
        if response.status_code == 200 or response.status_code == 206:
            # 创建下载目录
            os.makedirs("downloads", exist_ok=True)
            
            # 生成安全文件名
            safe_title = "".join(c for c in title if c.isalnum() or c in " -_")
            if not safe_title:
                safe_title = "douyin_video"
            
            filename = f"downloads/{safe_title}_{int(time.time())}.mp4"
            
            # 获取文件大小
            file_size = int(response.headers.get('Content-Length', 0))
            
            # 下载进度跟踪
            downloaded = 0
            chunk_size = 1024 * 1024  # 1MB
            
            with open(filename, 'wb') as f:
                for chunk in response.iter_content(chunk_size=chunk_size):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if file_size > 0:
                            percent = (downloaded / file_size) * 100
                            print(f"下载进度: {percent:.1f}%", end='\r')
            
            download_time = time.time() - start_time
            actual_size = os.path.getsize(filename) / (1024 * 1024)  # MB
            
            print(f"\n视频已保存至: {filename}")
            print(f"文件大小: {actual_size:.2f} MB")
            print(f"下载时间: {download_time:.2f} 秒")
            return filename
        else:
            print(f"下载失败，状态码: {response.status_code}")
            print(f"响应头: {response.headers}")
            return None
    except Exception as e:
        print(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return None
    finally:
        if driver:
            driver.quit()
            print("浏览器已关闭")

if __name__ == "__main__":
    # 抖音视频链接
    share_url = "https://v.douyin.com/X01lRooQLDc/"
    
    # 使用非无头模式以便处理验证
    result = download_douyin_video(share_url, headless=False)
    
    if result:
        print(f"成功下载视频: {result}")
    else:
        print("视频下载失败")